#include <iostream>
#include <string>
#include <chrono>
#include <winsock2.h>
#include <ws2tcpip.h>
#pragma comment(lib, "ws2_32.lib")

#define CLOSE_SOCKET closesocket

// ICMP 报文结构
struct icmp_header
{
    uint8_t type;
    uint8_t code;
    uint16_t checksum;
    uint16_t id;
    uint16_t seq;
};

// 计算校验和
uint16_t checksum(void *buf, int len)
{
    uint32_t sum = 0;
    auto *data = (uint16_t *)buf;
    while (len > 1)
    {
        sum += *data++;
        len -= 2;
    }
    if (len == 1)
    {
        sum += *(uint8_t *)data;
    }
    sum = (sum >> 16) + (sum & 0xFFFF);
    sum += (sum >> 16);
    return (uint16_t)~sum;
}

int main(int argc, char *argv[])
{
    if (argc != 2)
    {
        std::cerr << "Usage: " << argv[0] << " <host>" << std::endl;
        return 1;
    }

    // 初始化 Winsock
    WSADATA wsa;
    if (WSAStartup(MAKEWORD(2, 2), &wsa) != 0)
    {
        std::cerr << "WSAStartup failed" << std::endl;
        return 1;
    }

    // 创建原始套接字
    SOCKET sock = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
    if (sock == INVALID_SOCKET)
    {
        std::cerr << "Socket creation error: " << WSAGetLastError() << std::endl;
        WSACleanup();
        return 1;
    }

    // 设置超时
    DWORD timeout = 3000;
    if (setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, (const char *)&timeout, sizeof(timeout)) == SOCKET_ERROR)
    {
        std::cerr << "Setsockopt failed: " << WSAGetLastError() << std::endl;
        CLOSE_SOCKET(sock);
        WSACleanup();
        return 1;
    }

    // 解析目标地址
    sockaddr_in dest{};
    dest.sin_family = AF_INET;
    if (inet_addr(argv[1]) == INADDR_NONE)
    {
        std::cerr << "Invalid address" << std::endl;
        CLOSE_SOCKET(sock);
        WSACleanup();
        return 1;
    }
    dest.sin_addr.s_addr = inet_addr(argv[1]);

    // 构造 ICMP 报文
    char packet[sizeof(icmp_header) + 32];
    icmp_header *icmp = (icmp_header *)packet;
    icmp->type = 8; // Echo Request
    icmp->code = 0;
    icmp->checksum = 0;
    icmp->id = (uint16_t)GetCurrentProcessId();
    icmp->seq = 0;

    // 填充数据
    auto start = std::chrono::steady_clock::now();
    memset(packet + sizeof(icmp_header), 0, 32);

    // 计算校验和
    icmp->checksum = checksum(packet, sizeof(packet));

    // 发送报文
    if (sendto(sock, packet, sizeof(packet), 0, (sockaddr *)&dest, sizeof(dest)) == SOCKET_ERROR)
    {
        std::cerr << "Send failed: " << WSAGetLastError() << std::endl;
        CLOSE_SOCKET(sock);
        WSACleanup();
        return 1;
    }

    // 接收响应
    char buffer[1024];
    sockaddr_in from{};
    int fromlen = sizeof(from);
    int rcvd = recvfrom(sock, buffer, sizeof(buffer), 0, (sockaddr *)&from, &fromlen);

    if (rcvd == SOCKET_ERROR)
    {
        std::cerr << "Timeout or error: " << WSAGetLastError() << std::endl;
    }
    else
    {
        auto end = std::chrono::steady_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);

        // 解析响应
        icmp_header *ricmp = (icmp_header *)(buffer + 20); // 跳过 IP 头
        if (ricmp->type == 0 && ricmp->id == icmp->id)
        {
            char ipstr[INET_ADDRSTRLEN];
            strncpy(ipstr, inet_ntoa(from.sin_addr), INET_ADDRSTRLEN);
            std::cout << "Reply from " << ipstr
                      << ": bytes=" << rcvd - 20 - sizeof(icmp_header)
                      << " time=" << duration.count() << "ms"
                      << std::endl;
        }
    }

    CLOSE_SOCKET(sock);
    WSACleanup();
    return 0;
}